import random
import unittest
from Bio import SeqIO
from StringIO import StringIO
import task2

class TestFunctions(unittest.TestCase):
    def setUp(self):
        self.dna = ""
        self.rna = ""
        for i in xrange(12):
            t = random.randint(0,3)
            if t == 0:
                self.dna += "A"
                self.rna += "A"
            elif t == 1:
                self.dna += "C"
                self.rna += "C"
            elif t == 2:
                self.dna += "G"
                self.rna += "G"
            else:
                self.dna += "T"
                self.rna += "U"
        fasta_str = ">1\n" + self.dna + "\n"
        self.record = SeqIO.read(StringIO(fasta_str), "fasta")

    def test_to_rna(self):
        rna_to_test = task2.to_rna(self.record)
        self.assertEqual(rna_to_test._data, self.rna)

    def test_file_extension(self):
        self.assertTrue(task2.get_extension("input.fasta") == ".fasta")
        self.assertTrue(task2.get_extension("input_data.fa") == ".fa")
        self.assertTrue(task2.get_extension("../data/input.gbk") == ".gbk")
        self.assertTrue(task2.get_extension("../data/1/2/input.gb") == ".gb")
        self.assertRaises(Exception, task2.get_extension, "input.ext", msg="Unknown file format\t.ext")

    def test_file_exists(self):
        name = "abracadabra"
        with self.assertRaises(Exception):
            task2.get_records([name])


if __name__ == "__main__":
    unittest.main()


